package william.heap;

import william.utils.ArrayUtils;

import java.util.Arrays;

/**
 * @author ZhangShenao
 * @date 2024/2/19
 * @description 实现大根堆, 核心就是shiftUp和shiftDown两个方法
 */
public class MaxHeap {
    public static void main(String[] args) {
        int len = 10;
        int[] arr = ArrayUtils.generateRandomArray(len, 0, 999);
        System.out.println("原始数组: " + Arrays.toString(arr));

        MaxHeap maxHeap = new MaxHeap(len);
        for (int i = 0; i < len; i++) {
            maxHeap.offer(arr[i]);
        }

        System.out.println("依次获取堆顶元素: ");
        for (int i = 0; i < len; i++) {
            System.out.print(maxHeap.poll() + " ");
        }
        System.out.println();
    }

    /**
     * 底层采用数组维护元素
     */
    private int[] arr;

    /**
     * 堆的最大容量
     */
    private int capacity;

    /**
     * 记录堆中元素数量
     */
    private int size;

    /**
     * 构造器
     */
    public MaxHeap(int capacity) {
        this.arr = new int[capacity];
        this.capacity = capacity;
        this.size = 0;
    }

    /**
     * 获取堆中元素数量
     */
    public int size() {
        return size;
    }

    /**
     * 判断堆是否为空
     */
    public boolean isEmpty() {
        return (size == 0);
    }


    /**
     * 向堆中插入元素
     */
    public void offer(int value) {
        //校验长度
        if (size == capacity) {
            throw new UnsupportedOperationException("heap is full!!");
        }

        //将新元素插入数组末尾,并将堆长度+1
        arr[size++] = value;

        //进行shiftUp上浮操作
        shiftUp(size - 1);
    }

    /**
     * 查看堆顶元素
     */
    public int peek() {
        //校验长度
        if (isEmpty()) {
            throw new UnsupportedOperationException("heap is empty!!");
        }

        //直接返回根节点
        return arr[0];
    }

    /**
     * 获取并弹出堆顶元素
     */
    public int poll() {
        //校验长度
        if (isEmpty()) {
            throw new UnsupportedOperationException("heap is empty!!");
        }

        int root = arr[0];

        //将数组末尾元素赋值给根节点,并将size减1
        arr[0] = arr[--size];

        //进行shiftDown操作
        shiftDown(0);

        //返回根节点
        return root;
    }

    /**
     * 上浮操作
     * 将新插入的节点与父节点进行比较,如果比父节点大,则与父节点交换。重复该过程,直到到达根节点,或当前节点小于父节点
     */
    private void shiftUp(int idx) {
        while (idx > 0 && arr[idx] > arr[parentIdx(idx)]) {
            ArrayUtils.swap(arr, idx, parentIdx(idx));
            idx = parentIdx(idx);
        }
    }

    /**
     * 下沉操作
     * 从根节点开始,将该节点与它的左、右子节点中的最大值进行比较,如果最大值大于当前元素,则将该元素与最大的子节点进行交换
     * 重复上述过程,直到到达叶子节点,或当前节点大于左、右子节点
     */
    private void shiftDown(int idx) {
        while (idx < size) {
            int leftIdx = leftChildIdx(idx);
            if (leftIdx >= size) {   //无左子节点,终止操作
                break;
            }
            //找到左、右子节点中的最大节点的索引
            int maxIdx = leftIdx;
            int rightIdx = rightChildIdx(idx);
            if (rightIdx < size && arr[rightIdx] > arr[leftIdx]) {
                maxIdx = rightIdx;
            }


            if (maxIdx == idx) {     //当前父节点已经是最大值,终止操作
                break;
            }

            //将父节点与最大节点进行交换,并重复该过程
            ArrayUtils.swap(arr, idx, maxIdx);
            idx = maxIdx;
        }
    }

    /**
     * 获取父节点索引
     */
    private int parentIdx(int idx) {
        return (idx - 1) / 2;
    }

    /**
     * 获取左子节点索引
     */
    private int leftChildIdx(int idx) {
        return idx * 2 + 1;
    }

    /**
     * 获取右子节点索引
     */
    private int rightChildIdx(int idx) {
        return idx * 2 + 2;
    }
}
